#!/home/pathop/anaconda3/bin/python

import wrf,metpy
import numpy as np
import pandas as pd

from netCDF4 import Dataset

import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
plt.switch_backend('agg')

import metpy.calc as mpcalc
from metpy.cbook import get_test_data
from metpy.plots import Hodograph, SkewT
from metpy.units import units

#get files generation time
import time

import datetime
import os

from wrf_files_selected import wrfout_files, structure
#[structure]:date path 

out_dir='/public2/home/pathop/for_boanda_hubei/TSKEW_PLOT/'
"'read station data'"
station_file="/public2/home/pathop/boanda_hubei/stn_list_hubei.txt"
sta=np.loadtxt(station_file,dtype=bytes)

#get station name
sta_name=sta[:,1].astype(str)
#get longititude
sta_lon=sta[:,2].astype(float)
#get latititude
sta_lat=sta[:,3].astype(float)
#get number of station
n_station=sta.shape[0]


"'allocate vertical levels for Skew-T plot'"
#底层层数调整，避免skew图底层风向太密集。
start=0
end=38
step=4
bottom=step*3
interm=38-16
index=np.arange(start,bottom,step).tolist()+np.arange(bottom,interm,2).tolist()+np.arange(interm,end,1).tolist()
#index用于选取垂直方向数据

def set_skew(pressure, temperature,dewpoint,u_wind, v_wind,prefix_out_jpg,name,year,month,day,hour,minu):
    # Create a new figure. The dimensions here give a good aspect ratio
    fig = plt.figure(figsize=(8, 9))
    #add_metpy_logo(fig, 630, 80, size='large')

    # Grid for plots
    gs = gridspec.GridSpec(3, 3)
    rotation=50
    skew = SkewT(fig, subplot=gs[:, :2],rotation=rotation)#,rect=[0,0,0.9,1])   #rotation:Controls the rotation of temperature relative to horizontal

    # Plot the data using normal plotting functions, in this case using
    # log scaling in Y, as dictated by the typical meteorological plot
    temp_line=skew.plot(pressure, temperature, 'r',label='temp')
    dew_line=skew.plot(pressure, dewpoint, 'g')
    skew.plot_barbs(pressure, u_wind, v_wind)
    skew.ax.set_ylim(1025, 100)
#    skew.ax.legend((temp_line,dew_line),('temperature','dewpoint'))
    #skew.axes

    ############################################
    #字体设置
    fontsize=15
    ax=plt.gca()
    for tick in ax.xaxis.get_major_ticks():
         tick.label.set_fontsize(fontsize)
    for tick in ax.yaxis.get_major_ticks():
         tick.label.set_fontsize(fontsize)

    #ax.legend((temp_line,dew_line),('temperature','dewpoint'))
    #define ticks
    skew.ax.set_yticks([1000,850,700,500,400,300,250,200,150,100])
    #define xaxis label
    skew.ax.set_xlabel('temperature (degC)',fontsize=15)
    #define yaxis label
    skew.ax.set_ylabel('pressure (hPa)',fontsize=15)
    #set figure title()
    the_beijing_time=year+'-'+month+'-'+day+' '+hour+':'+minu
    skew.ax.set_title('Skew-T Plot \nStation: {} Time: {}'.format(name,\
                      the_beijing_time), fontsize=15, loc='left')

    ############################################
    # Add the relevant special lines
    dry_adiabats=skew.plot_dry_adiabats(linewidth=0.8,linestyle='-.') #干绝热线
    moist_adiabat=skew.plot_moist_adiabats(linewidth=0.8,linestyle='-.') #湿绝热线
    mixing_line=skew.plot_mixing_lines(linewidth=0.8,linestyle='-.') #混合线

    # Good bounds for aspect ratio
    biggest=temperature.max()
    length=calcu_length(rotation,T=temperature,Td=dewpoint,press=pressure)
    smallest=biggest-length
    max_num=biggest+5
    min_num=smallest-5
    skew.ax.set_xlim(min_num, max_num )

    #set a place to lay out the legend
    ax1 = fig.add_subplot(gs[1, -1])
    # Create a hodograph
    ax2 = fig.add_subplot(gs[0, -1])
    h = Hodograph(ax2, component_range=60.)  #
    h.add_grid(increment=20)
    h.plot( u_wind, v_wind)

    #add figure legend
    legend=ax1.legend([temp_line[0],dew_line[0],mixing_line,moist_adiabat,dry_adiabats],\
                      ["temperature","dew point","mixing line","moist adiabat","dry adiabats"],\
                      fontsize=15,loc=2, ncol=1)#, mode='expand')

    legend.get_frame().set_edgecolor('none')
    ax1.axis('off')


    out_jpg=prefix_out_jpg+'_'+name+'.jpg'

    plt.savefig(out_jpg,dpi=60)
    print('a plot has been saved')
    plt.cla()
    plt.close("all")
    del fig
    del skew
    return

def calcu_length(rotation,T,Td,press):
    import math
    cut=T.max()-Td #distance betwen the max T and the min Td, which is an array.
    rotat=rotation*math.pi/180 #convert angle to radian
    #use [rotat] as the angle to caculate
    "'get top pressure'"
    index=str(press.min()>100)
    option={'True': 100,'False': press.min()}
    top_press=option[index]
    whole_height=press.max()-top_press #[top_press]: the min num of pressure
    similarity_prop=(press.max()-press)/whole_height
    #similar proportions, which is an array.
    deduction=cut*((math.sin(rotat)**2))*similarity_prop
    #[deduction] is means the excised part of axis, and [deduction] is an array.
    length=(cut-deduction).max()
    #依据：cut(1-sin**2)=cut*cos**2
    #get the max num
    return length



def main(out_dir):
    #global rh,tc,tdc,height,u,v,press,td
    for wrfout_afile in wrfout_files:
        wrfout=Dataset(wrfout_afile)
        "'get the location which corresponding to the station'"
        xx,yy=wrf.ll_to_xy(wrfin=wrfout,latitude=sta_lat,longitude=sta_lon)
        #print("xx:",xx.values)
        #print('yy:',yy.values)
        "'recognition time'"
        Times=wrfout.variables['Times'][:]
        ntimes=Times.shape[0]
        out_time=wrf.getvar(wrfout,'times').values
        "'clip wrf time'"
        if ntimes==1:
            tim=str(out_time)
            year=tim[:4]#int(tim[:4])
            month=tim[5:7]#int(tim[5:7])
            day=tim[8:10]#int(tim[8:9])
            UTC_hour=tim[11:13]#int(tim[11:13])
            #hour='%02d'%(int(UTC_hour)+8)
            minu=tim[14:16]
            #pd.DatetimeIndex(out_time)

        UTC_time=datetime.datetime(year=int(year),month=int(month),day=int(day),hour=int(UTC_hour),minute=int(minu))

        delta=datetime.timedelta(hours=8)

        beijing_time=UTC_time+delta

        year='%04d'%beijing_time.year
        month='%02d'%beijing_time.month
        day='%02d'%beijing_time.day
        hour='%02d'%beijing_time.hour
        minu='%02d'%beijing_time.minute

        year_UTC='%04d'%UTC_time.year
        month_UTC='%02d'%UTC_time.month
        day_UTC='%02d'%UTC_time.day
        hour_UTC='%02d'%UTC_time.hour
        minu_UTC='%02d'%UTC_time.minute

        #appoint the output directory
        out_path=out_dir+structure
        #out_path=out_dir+'/'+year+'/'+year+month+'/'+year+month+day+hour+'/'

        "'Judging whether the out_path exist'"
        if os.path.isdir(out_path):
            print('path is there')
        else:
            os.makedirs(out_path)
            print('path is there create')

        "'getting wrf variables from output,these variables are 3-d[vertical,lat,lon] strucured frame data'"
        rh=  wrf.getvar(wrfout,"rh")
        tc=wrf.getvar(wrfout,"tc")
        tdc=wrf.getvar(wrfout,"td")
        height=wrf.getvar(wrfout,"height")
        u= wrf.getvar(wrfout,"ua")   # u on mass points
        v= wrf.getvar(wrfout,"va")   # v on mass points
        press=  wrf.getvar(wrfout,"p")*0.01
        td=  wrf.getvar(wrfout,"td")
        "'getting station location reference from wrfout data '"
        height_loc=height[:,yy,xx].values
        press_loc=press[:,yy,xx].values
        tc_loc=tc[:,yy,xx].values
        td_loc=td[:,yy,xx].values
        u_loc=u[:,yy,xx].values
        v_loc=v[:,yy,xx].values
        "'delete redundant variables'"
        del rh,tc,tdc,height,u,v,press,td
        "'set output files'"
        prefix_out_jpg=out_path+"vertical_t_dew_"+year_UTC+month_UTC+day_UTC+hour_UTC
        "'Date & Time detection'"
        #print(prefix_out_jpg)
        #print('UTC_time:',UTC_time)
        #print('beijing time:',beijing_time)
        "'calling function to plot'"
#############################################################
        for order in range(n_station):
            
            #get the station name
            name=str(sta_name[order])
            print('the station order is:',order)
            print('the station name is:',name)
            "'define DataFrame'"
            col_names = ['pressure', 'height', 'temperature', 'dewpoint', 'direction', 'speed','u_wind', 'v_wind']
            df = pd.DataFrame(columns = col_names)
           # print(p[:,:,:].values.shape)
           # print('\n\n\n')
           # print('yy:',yy)
           # print('xx:',xx)
            
            #.values[level,num]中的num可以是1~13，代表不同站点
            df['height_']=height_loc[:,order]
            df['pressure_']=press_loc[:,order]
            df['temperature_']=tc_loc[:,order]
            df['dewpoint_']=td_loc[:,order]
            df['u_wind_']=u_loc[:,order]
            df['v_wind_']=v_loc[:,order]
            
            pressure=df['pressure_'].values[index]
            temperature=df['temperature_'].values[index]

            dewpoint=df['dewpoint_'].values[index]
            u_wind=df['u_wind_'].values[index]
            v_wind=df['v_wind_'].values[index]
            
#############################################################
            set_skew(pressure, temperature,dewpoint,u_wind, v_wind,prefix_out_jpg,name,year,month,day,hour,minu)
        
def text_create(afile_name,afile_path):
    "'afile_path:save path of the new file'"
    full_path = afile_path + afile_name + '.txt'  # 也可以创建一个.doc的word文档
    file = open(full_path, 'w')
    msg="not found input files at time:"+afile_name
    file.write(msg)
    file.close()

def error_message():
    "'Specify the output path'"
    
    "'Estimate whether the file exist'"
    target=str(wrfout_files==[])
    afile_name=os.popen('date +"%Y-%m-%d_%H:%M.%S"').read()
    situation={'True': text_create(afile_name,afile_path=out_dir+structure),
               'False':True}
    situation[target]
    return

    
if __name__ == '__main__':
    global rh,tc,tdc,height,u,v,press,td
    error_message()
    main(out_dir)
    print("run complete!")

